import os from typing import Dict, List import uuid import csv import numpy as np from PIL import Image def write_row(filepath:str, row: Dict): new_file = not os.path.isfile(filepath) with open(filepath, mode="a", newline="") as file: fieldnames = row.keys() writer = csv.DictWriter(file, fieldnames=fieldnames) if new_file: writer.writeheader() # Write header if new file writer.writerow(row) # Write the row class Feedback(): def __init__(self, image_dir = './data/input', mask_dir = './data/mask', inference_csv = './data/inference.csv', feedback_csv = './data/feedback.csv', ): self.image_dir = image_dir self.mask_dir = mask_dir self.inference_csv = inference_csv self.feedback_csv = feedback_csv def save_inference(self, pt_coords:List, pt_labels:List, image: Image.Image, mask: np.ndarray): self.inference_id = uuid.uuid4() write_row( filepath=self.inference_csv, row = { "inference_id": self.inference_id, "image": image.tobytes(), "mask": mask.tobytes(), "pt_coords": str(pt_coords), "pt_labels": str(pt_labels), } ) def save_feedback(self, cutout_idx:int=None, feedback_str:str=None, like:int=None): write_row( filepath=self.feedback_csv, row = { "inference_id": self.inference_id, "cutout_idx": cutout_idx, "feedback_str": feedback_str, "like": like, } )