dor6_vision/main_eval_eda.py at master · amitli1/dor6_vision · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
from tqdm import tqdm
import pandas as pd
import glob
import os
import torch

import plotly.express as px
import matplotlib.pyplot as plt

from app_config.settings import TRAIN_FULL_MODE_FILES_PATH, TRAIN_CROP_FILES
from main_create_embeddings import create_embeddings, load_model, run_umap
from pointing_agent.pointing_agent import PointingAgent
from PIL import Image

def get_testset_pointing(test_set_path, output_name):
    pointingAgent = PointingAgent()

    l_jpg_files = glob.glob(f"{test_set_path}/*.jpg")

    l_jpg = []
    l_point_pred = []
    l_x = []
    l_y = []

    for full_file_path in tqdm(l_jpg_files):
        file                 = os.path.basename(full_file_path)
        generated_text, x, y = pointingAgent.run_molmo_prediction(full_file_path)

        l_jpg.append(file)
        l_point_pred.append(generated_text)
        l_x.append(x)
        l_y.append(y)

    df_res = pd.DataFrame({'jpg_file': l_jpg,"point_pred": l_point_pred, "x": l_x, "y": l_y})
    df_res.to_csv(output_name, index=False)

    return df_res


def crop_images(image_path, x, y, output_crop_path):

    base_name      = os.path.basename(image_path)
    image          = Image.open(image_path).convert("RGB")
    width, height  = image.size

    crop_size = 200
    left      = max(0, x - crop_size / 2)
    top       = max(0, y - crop_size / 2)
    right     = min(width, x + crop_size / 2)
    bottom    = min(height, y + crop_size / 2)

    crop_image = image.crop((left, top, right, bottom))
    # crop_image.show()
    crop_image.save(f"{output_crop_path}/{base_name}", "JPEG")

def crop_test_files(testset_with_point_csv, output_crop_path):

    df_point = pd.read_csv(testset_with_point_csv)
    for i in tqdm(range(len(df_point))):
        image_path = f"/home/amitli/repo/dor6_vision/Dataset/test_set/{df_point.jpg_file.values[i]}"
        crop_images(image_path, df_point.x.values[i], df_point.y.values[i], output_crop_path)

def test_set_crop_embeddings():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    processor, model = load_model(f"facebook/dinov2-base", device)
    create_embeddings(model, processor, device,
                      TEST_SET_POINT_CSV,
                      TEST_SET_CROP_FOLDER,
                      "/home/amitli/repo/dor6_vision/Dataset/embeddings_testset_crop.csv")

def run_umap_train_and_test_crop():
    df_train           = pd.read_csv("/home/amitli/repo/dor6_vision/Dataset/embeddings_train_crop.csv")
    df_test            = pd.read_csv("/home/amitli/repo/dor6_vision/Dataset/embeddings_testset_crop.csv")
    df_train_test_crop = pd.concat([df_train, df_test], ignore_index=True)
    df_train_test_crop.to_csv("/home/amitli/repo/dor6_vision/Dataset/train_test_crop_embeddings.csv")
    run_umap("/home/amitli/repo/dor6_vision/Dataset/train_test_crop_embeddings.csv",
             "/home/amitli/repo/dor6_vision/Dataset/train_test_umap_crop.csv")


def plot_umap_df(df_file_name):
    df        = pd.read_csv(df_file_name)
    color_map = {
        'SA-22'  : 'red',
        'SCUD'   : 'green',
        'T-90'   : 'blue',
        'unknown': 'black',
    }

    # Map colors to df['gt']
    colors = df['gt'].map(color_map)

    plt.figure(figsize=(8, 6))
    plt.scatter(df['umap_x'], df['umap_y'], c=colors, s=40, alpha=0.8)

    plt.xlabel('UMAP X')
    plt.ylabel('UMAP Y')
    plt.title(f'UMAP projection colored by GT ({len(df)} Samples)')
    plt.grid(True)

    # Add legend
    for cls, color in color_map.items():
        plt.scatter([], [], c=color, label=cls)
    plt.legend(title='GT')

    plt.show()


def plot_umap_to_html(df_file_name, plot_file_name):
    df          = pd.read_csv(df_file_name)
    df['hover'] = df["jpg_file"]
    fig = px.scatter(df, x='umap_x', y='umap_y', color='gt', hover_name="hover", title=f"#Samples: {len(df)}")

    fig.write_html(plot_file_name)
    fig.show()

def print_statisics():

    full_file = f"{TRAIN_FULL_MODE_FILES_PATH}/11-17-44_444400_23.jpg"
    crop_file = f"{TRAIN_CROP_FILES}/11-17-44_444400_23.jpg"
    full_image = Image.open(full_file).convert("RGB")
    crop_image = Image.open(crop_file).convert("RGB")
    print(f"Full file size:{full_image.size}")
    print(f"Crop file size:{crop_image.size}")

if __name__ == "__main__":

    # print_statisics()
    # exit(0)

    TRAIN_EMBEDDINGS_CROP_CSV = "/home/amitli/repo/dor6_vision/Dataset/embeddings_crop.csv"
    TEST_SET_PATH             = "/home/amitli/repo/dor6_vision/Dataset/test_set/"
    TEST_SET_CROP_FOLDER      = "/home/amitli/repo/dor6_vision/Dataset/test_set_crop/"
    TEST_SET_POINT_CSV        = "/home/amitli/repo/dor6_vision/Dataset/test_set_point.csv"

    #get_testset_pointing(TEST_SET_PATH, TEST_SET_POINT_CSV)
    #crop_test_files(TEST_SET_POINT_CSV, TEST_SET_CROP_FOLDER)

    #test_set_crop_embeddings()
    run_umap_train_and_test_crop()

    plot_umap_df("/home/amitli/repo/dor6_vision/Dataset/train_test_umap_crop.csv")
    plot_umap_to_html("/home/amitli/repo/dor6_vision/Dataset/train_test_umap_crop.csv",
                      "/home/amitli/repo/dor6_vision/Dataset/train_test_crop.html")