pypoint.py

from workspace import storage
import argparse
import numpy as np
import cv2
import os
import pickle
import json # Added for COCO format saving
import config

"""TODO
        1. Add previous button
        2. Load labels from saved data on displayed image
        3. Add b.b labeling feature
        4. Add MS COCO Saving format
        5. Add save as json
"""

# ----------------- Define cml arguments ---------------------- #
parser = argparse.ArgumentParser()
parser.add_argument('--dir', type=str, required=True, default='single', help='The path to the file in which the images exist.')
parser.add_argument('--classes',type=str, required=True, help='The name of the classes to be labeled. A list of str.')
parser.add_argument('--workspace',type=str, required=False, default='./', help='Specifies a directory to save the last labeled item info.')
parser.add_argument('--savedir',type=str, required=False, default='./', help='Specifies a directory to save the labeled data.')

# ------------- Parse cml arguments and set config ------------------ #
args = parser.parse_args()
config.dir = args.dir
config.save_dir = args.savedir
config.workspace = args.workspace
config.classes = args.classes.strip('[]').split(',')
config.set_colors(len(config.classes))

def update_stat():
    """
        Updates the status bar context
    """
    config.img_stat[img.shape[0]:, :,:] = 0
    font = cv2.FONT_HERSHEY_SIMPLEX
    # Display current class and labeling mode
    mode_text = f"Mode: {config.labeling_mode.capitalize()}"
    cv2.putText(config.img_stat,f"Labeling {config.classes[config.curr_class_idx]} ({mode_text})", (10, img.shape[0]+30), font, 1, (255,255,255),2)
    # Update guide text
    guide_text = "Guide: q:exit / n:save&next / p:prev / u:undo / m:mode / j:COCO / k:JSON / LMouse:label / RMouse:next_cls"
    if config.labeling_mode == 'bbox':
        guide_text = "Guide: q:exit / n:save&next / p:prev / u:undo / m:mode / j:COCO / k:JSON / LMouse:bbox / RMouse:next_cls"
    cv2.putText(config.img_stat, guide_text, (10, img.shape[0]+60), font, 0.5, (255,255,255),1)
    
    # Console message (optional, can be removed or refined)
    # print(f"Class: {config.classes[config.curr_class_idx]}, Mode: {config.labeling_mode}")


# Function to save labels in MS COCO format for the current image
def save_coco_format(image_filename, image_shape, labels_data, classes_list, save_directory):
    """
    Saves the labels for the current image in MS COCO JSON format.
    image_filename: string, the name of the image file
    image_shape: tuple, (height, width, channels) of the image
    labels_data: dict, the current labels in the format {'class_name': [{'type': 'point'/'bbox', 'coords': ...}]}
    classes_list: list, list of all available class names
    save_directory: string, path to the directory where the COCO file should be saved
    """
    import datetime # For timestamp

    current_time = datetime.datetime.now()

    coco_output = {
        "info": {
            "description": "Labels generated by pypoint tool",
            "version": "1.0",
            "year": current_time.year,
            "contributor": "pypoint user",
            "date_created": current_time.isoformat()
        },
        "licenses": [
            {
                "id": 1,
                "name": "Placeholder License",
                "url": ""
            }
        ],
        "images": [],
        "annotations": [],
        "categories": []
    }

    # Populate categories
    for idx, class_name in enumerate(classes_list):
        coco_output["categories"].append({
            "id": idx, # Using list index as category ID
            "name": class_name,
            "supercategory": class_name # Or a more general one if applicable
        })

    # Populate image information
    height, width, _ = image_shape
    coco_output["images"].append({
        "id": 0, # For single-image COCO export, image_id is 0
        "width": int(width),
        "height": int(height),
        "file_name": image_filename,
        "license": 1, # Reference to the placeholder license
        "date_captured": "" # Optional: could be file modification time
    })

    # Populate annotations
    annotation_id_counter = 0
    for class_name, items_list in labels_data.items():
        try:
            # Find the category_id corresponding to class_name
            category_id = next(cat['id'] for cat in coco_output['categories'] if cat['name'] == class_name)
        except StopIteration:
            print(f"Warning: Class '{class_name}' not found in COCO categories. Skipping its annotations.")
            continue

        for item in items_list:
            ann_obj = {
                'id': annotation_id_counter,
                'image_id': 0, # Corresponds to the single image ID in coco_output['images']
                'category_id': category_id,
                'iscrowd': 0
            }
            item_type = item.get('type')
            coords = item.get('coords')

            if item_type == 'point' and coords:
                x, y = coords
                ann_obj['segmentation'] = [[float(x), float(y)]]
                ann_obj['bbox'] = [float(x), float(y), 1.0, 1.0] # COCO bbox: [x,y,width,height]
                ann_obj['area'] = 1.0
            elif item_type == 'bbox' and coords:
                x1, y1, x2, y2 = coords
                # Ensure x1<x2, y1<y2 was handled during creation, but good practice for COCO:
                width = float(x2 - x1)
                height = float(y2 - y1)
                coco_bbox = [float(x1), float(y1), width, height]
                
                ann_obj['bbox'] = coco_bbox
                # Polygon format for bbox segmentation:
                ann_obj['segmentation'] = [[float(x1), float(y1), float(x1), float(y2), float(x2), float(y2), float(x2), float(y1)]]
                ann_obj['area'] = width * height
            else:
                print(f"Warning: Skipping unknown or malformed label item: {item}")
                continue # Skip this annotation
            
            coco_output["annotations"].append(ann_obj)
            annotation_id_counter += 1

    # Save to JSON file
    # Ensure save_directory exists (optional, good practice)
    # if not os.path.exists(save_directory):
    #     os.makedirs(save_directory)

    base_filename, _ = os.path.splitext(image_filename)
    output_filename = os.path.join(save_directory, f'{base_filename}.coco.json')
    
    try:
        with open(output_filename, 'w') as f:
            json.dump(coco_output, f, indent=4)
        print(f"Labels saved in COCO format to: {output_filename}")
    except IOError as e:
        print(f"IOError saving COCO JSON: {e}")
    except Exception as e:
        print(f"An unexpected error occurred while saving COCO JSON: {e}")

# Function to save labels in a generic JSON format for the current image
def save_generic_json(image_filename, image_shape, labels_data, save_directory):
    """
    Saves the labels for the current image in a generic JSON format.
    image_filename: string, the name of the image file
    image_shape: tuple, (height, width, channels) of the image
    labels_data: dict, the current labels in the format {'class_name': [{'type': 'point'/'bbox', 'coords': ...}]}
    save_directory: string, path to the directory where the JSON file should be saved
    """
    height, width, _ = image_shape
    
    generic_output = {
        'image_filename': image_filename,
        'image_dimensions': {
            'width': int(width),
            'height': int(height)
        },
        'labels': labels_data # Directly use the existing structure
    }

    base_filename, _ = os.path.splitext(image_filename)
    output_filename = os.path.join(save_directory, f'{base_filename}.generic.json')

    try:
        with open(output_filename, 'w') as f:
            json.dump(generic_output, f, indent=4)
        print(f"Labels saved in generic JSON format to: {output_filename}")
    except IOError as e:
        print(f"IOError saving generic JSON: {e}")
    except Exception as e:
        print(f"An unexpected error occurred while saving generic JSON: {e}")


def click_event(event, x, y, flags, params):
    """
        Handles the mouse click event
    """
 
    # checking for left mouse clicks
    if event == cv2.EVENT_LBUTTONDOWN:
        
        if config.undo:
            config.undo = False
        # displaying the coordinates
        # on the Shell
        print(x, ' ', y)
        
        current_class_name = config.classes[config.curr_class_idx]
        color = config.colors[config.curr_class_idx]

        # Ensure the list for the current class exists
        if current_class_name not in config.curr_labels:
            config.curr_labels[current_class_name] = []

        config.last_image_cache = config.img_stat.copy()

        if config.labeling_mode == 'point':
            label_item = {'type': 'point', 'coords': (x, y)}
            config.curr_labels[current_class_name].append(label_item)
            cv2.circle(config.img_stat, (x, y), radius=4, color=color, thickness=-1)
            print(f"Point added: {label_item} for class {current_class_name}")

        elif config.labeling_mode == 'bbox':
            if config.bbox_start_point is None:
                # First click for bbox
                config.bbox_start_point = (x, y)
                # Optionally draw a temporary marker for the start point
                cv2.circle(config.img_stat, (x, y), radius=3, color=color, thickness=1) 
                print(f"Bbox start point: {config.bbox_start_point} for class {current_class_name}")
            else:
                # Second click for bbox
                x1, y1 = config.bbox_start_point
                x2, y2 = x, y

                # Ensure x1 < x2 and y1 < y2
                if x1 > x2: x1, x2 = x2, x1
                if y1 > y2: y1, y2 = y2, y1
                
                # Clear the temporary start point marker by restoring the pre-click image for this spot
                # This is a bit tricky if other labels are already there.
                # A simpler way is to just draw over it, or ensure last_image_cache is taken before the first bbox click marker.
                # For now, the start marker is small and will be part of the bbox line or filled by it.

                label_item = {'type': 'bbox', 'coords': (x1, y1, x2, y2)}
                config.curr_labels[current_class_name].append(label_item)
                cv2.rectangle(config.img_stat, (x1, y1), (x2, y2), color, thickness=2)
                config.bbox_start_point = None # Reset for next bbox
                print(f"Bbox added: {label_item} for class {current_class_name}")

        cv2.imshow('image', config.img_stat)
 
    # checking for right mouse clicks    
    elif event==cv2.EVENT_RBUTTONDOWN:
        if config.undo:
            config.undo = False
        
        config.bbox_start_point = None # Reset bbox if switching class mid-drawing
        set_curr_class_idx()
        if config.curr_class_idx> (len(config.classes) - 1):
            set_curr_class_idx(1)
        
        update_stat()
        cv2.imshow('image', config.img_stat)

    elif event== cv2.EVENT_MBUTTONDOWN:
        # undo operation
        current_class_name = config.classes[config.curr_class_idx]
        try:
            if not config.undo:
                if config.bbox_start_point is not None:
                    # Cancel an in-progress bbox drawing (first point was clicked)
                    # last_image_cache should be from *before* the temporary bbox start point marker was drawn
                    # This means last_image_cache must be updated BEFORE drawing the temp marker in LBUTTONDOWN.
                    # Let's assume it was.
                    config.bbox_start_point = None
                    # If a temporary marker was drawn, we need to restore the image state before that marker.
                    # This requires careful handling of last_image_cache.
                    # Current LBUTTONDOWN sets last_image_cache, then draws temp marker. So this is correct.
                    config.img_stat = config.last_image_cache 
                    cv2.imshow('image', config.img_stat)
                    print("Bbox drawing cancelled.")
                    # No actual label was added to config.curr_labels, so don't set config.undo = True
                    # and don't try to pop.
                    return # Exit event handler

                if current_class_name in config.curr_labels and config.curr_labels[current_class_name]:
                    config.curr_labels[current_class_name].pop()
                    # last_image_cache was set before the last label was drawn.
                    config.img_stat = config.last_image_cache 
                    cv2.imshow('image', config.img_stat)
                    config.undo = True
                    print("Last label removed.")
                else:
                    print("No labels to undo for the current class.")
            else:
                print("Only one step undo is supported. You can restart labeling process for this sample with the 'r' key.")
        except IndexError: 
            print("No labels to undo (IndexError).") # Should be caught by the check above
        except Exception as e:
            print(f"Error during undo: {e}")
        
    else:
        pass

def set_curr_class_idx(reset=-1):
    if reset==1:
        config.curr_class_idx = 0
    else:
        config.curr_class_idx +=1

def process_previous_action(current_idx):
    """
    Handles the logic for the 'previous' action.
    Modifies config state (curr_labels, bbox_start_point, curr_class_idx).
    Returns the new image index.
    """
    current_idx -= 1
    if current_idx < 0:
        current_idx = 0
    
    config.curr_labels = {}  # Clear labels
    config.bbox_start_point = None # Reset bbox state
    set_curr_class_idx(1)  # Reset to the first class
    print(f"Moved to previous image index: {current_idx}. Labels cleared.") # For feedback
    return current_idx

if __name__=="__main__":

    strg = storage.Storage('./workspace/')
    last_idx, last_fname = [0, ""]
    last_idx, last_fname = strg.load()
    files = os.listdir(config.dir)
    i = 0
    config.curr_labels = {} # Initialize/clear current labels before starting

    while i < len(files):
        if last_idx>0:
            # This logic is to resume from where the user left off.
            # We should clear curr_labels here too before potentially loading new ones.
            # However, the 'n' and 'p' handlers already clear config.curr_labels.
            # The global clear above handles the very first image load.
            # For resuming, the image will be skipped, then 'i' incremented,
            # and on the *actual* first image to be labeled, curr_labels will be empty.
            if i<=last_idx:
                i+=1
                continue
        # read the image
        img = cv2.imread(os.path.join(config.dir, files[i]), 1)
        config.img_stat = np.zeros((img.shape[0]+config.stat_height, img.shape[1], img.shape[2]), dtype=img.dtype)
        config.img_stat[:img.shape[0],:img.shape[1],:] = img

        # Load and draw existing labels if any
        label_file_path = os.path.join(config.save_dir, f'{files[i]}.pkl')
        if os.path.exists(label_file_path):
            try:
                with open(label_file_path, 'rb') as f:
                    loaded_labels = pickle.load(f)
                    config.curr_labels = loaded_labels # Overwrite with loaded labels
                    for class_name, label_items in config.curr_labels.items():
                        if class_name in config.classes: # Ensure class still exists
                            class_idx = config.classes.index(class_name)
                            color = config.colors[class_idx]
                            for item in label_items:
                                if item['type'] == 'point':
                                    x_coord, y_coord = item['coords']
                                    cv2.circle(config.img_stat, (x_coord, y_coord), radius=4, color=color, thickness=-1)
                                elif item['type'] == 'bbox':
                                    x1, y1, x2, y2 = item['coords']
                                    cv2.rectangle(config.img_stat, (x1, y1), (x2, y2), color, thickness=2)
                        else:
                            print(f"Warning: Class '{class_name}' from label file not in current config. Skipping.")
            except Exception as e:
                print(f"Error loading or drawing labels for {files[i]}: {e}")
                config.curr_labels = {} # Clear labels if loading failed

        # display the image
        update_stat()
        cv2.imshow('image', config.img_stat)    

        # Set mouse click event handler
        cv2.setMouseCallback('image', click_event)

        # listen for keys
        key = cv2.waitKey(0)
        if key == ord('q'):
            #Exit
            #storage.save(i, file)
            break
        elif key == ord('n'):
            # Save and proceed to the next file
            with open(os.path.join(config.save_dir, f'./{files[i]}.pkl'), 'wb') as f:
                pickle.dump(config.curr_labels, f)
                print(f"The label for {files[i]} was succesfully saved!")
                set_curr_class_idx(1)
            strg.save(i, files[i])
            config.curr_labels = {}
            i+=1
        
        elif key == ord('p'):
            # Go to the previous image
            i = process_previous_action(i)
            # The loop will now continue to the next iteration,
            # and since 'i' has been modified by process_previous_action,
            # it will load the previous (or first) image's data
            # and call update_stat() at the beginning of the loop.
            continue

        elif key == ord('m'): # Mode switch
            if config.labeling_mode == 'point':
                config.labeling_mode = 'bbox'
            else:
                config.labeling_mode = 'point'
            config.bbox_start_point = None # Reset bbox state on mode switch
            update_stat()
            cv2.imshow('image', config.img_stat)

        elif key == ord('j'): # Save current image's labels in COCO format
            if config.curr_labels:
                # Ensure img and files[i] are available from the main loop scope
                save_coco_format(files[i], img.shape, config.curr_labels, config.classes, config.save_dir)
            else:
                print("No labels to save in COCO format for the current image.")

        elif key == ord('k'): # Save current image's labels in generic JSON format
            if config.curr_labels:
                # Ensure img and files[i] are available from the main loop scope
                save_generic_json(files[i], img.shape, config.curr_labels, config.save_dir)
            else:
                print("No labels to save in generic JSON format for the current image.")

        elif key== ord('c'):
            # Clears the storage?
            pass
        elif key==ord('r'):
            config.curr_labels = {}
            config.bbox_start_point = None # Reset bbox state if clearing labels for current image


    # close the window
    cv2.destroyAllWindows()