process_photos.py

import os
import sys
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QLineEdit, 
    QPushButton, QHBoxLayout, QGridLayout, QScrollArea
)
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from PIL import Image
from ultralytics import YOLO
from PIL.ExifTags import TAGS
from datetime import datetime
import shutil
import re
import random


#CONSTANTS
IMAGE_PATH = "samplePhotos"
OUTPUT_PATH = "out"
SPLIT_INTO_GROUPS = True
TIME_DIFFERENCE = 5
COL_NUMS = 2
WINDOW_WIDTH = 1500
WINDOW_HEIGHT = 850
MODEL_PATH = "runs/detect/train2/weights/best.pt"

class SplitImagesByTimestamp():
    def __init__(self, photosPath, outPath, timeStampDifference):
        self.photosPath = photosPath
        self.outPath = outPath
        self.timeDif = timeStampDifference
    
    def get_photo_timestamp(self, photo_path):
        try:
            image = Image.open(photo_path)
            exif_data = image._getexif()
            if exif_data:
                for tag, value in exif_data.items():
                    decoded_tag = TAGS.get(tag, tag)
                    if decoded_tag == "DateTimeOriginal":
                        return datetime.strptime(value, "%Y:%m:%d %H:%M:%S")
        except Exception as e:
            print(f"Error reading {photo_path}: {e}")
        return None

    def group_photos_by_time(self):
        photos = []
        
        # Collect photos and their timestamps
        for filename in os.listdir(self.photosPath):
            if filename.lower().endswith('.jpg'):
                photo_path = os.path.join(self.photosPath, filename)
                timestamp = self.get_photo_timestamp(photo_path)
                if timestamp:
                    photos.append((photo_path, timestamp))
        
        # Sort photos by timestamp
        photos.sort(key=lambda x: x[1])
        
        # Group photos based on the time threshold
        groups = []
        current_group = []
        previous_timestamp = None

        for photo_path, timestamp in photos:
            if not previous_timestamp or (timestamp - previous_timestamp).total_seconds() > self.timeDif:
                if current_group:
                    groups.append(current_group)
                current_group = [photo_path]
            else:
                current_group.append(photo_path)
            previous_timestamp = timestamp
        
        # Add the last group
        if current_group:
            groups.append(current_group)

        return groups

    def create_output_folders(self, groups):
        if not os.path.exists(self.outPath):
            os.makedirs(self.outPath)
        
        for idx, group in enumerate(groups):
            group_folder = os.path.join(self.outPath, f"group_{idx + 1}")
            os.makedirs(group_folder, exist_ok=True)
            
            for photo_path in group:
                shutil.copy(photo_path, group_folder)

class ImageProcessorThread(QThread):
    image_processed = pyqtSignal(int, Image.Image, tuple)  # Signal to send processed image and bbox
    processing_canceled = pyqtSignal()  # Signal when processing is canceled

    def __init__(self, image_files, model, parent=None):
        super().__init__(parent)
        self.image_files = image_files
        self.model = model
        self.cancel_flag = False

    def run(self):
        for idx, image_path in enumerate(self.image_files):
            if self.cancel_flag:
                self.processing_canceled.emit()
                break

            cropped_image, bbox = self.process_image(image_path)
            if cropped_image:
                # If no bbox was found, use a tuple of zeros
                if bbox is None:
                    bbox = (0, 0, 0, 0)
                self.image_processed.emit(idx, cropped_image, bbox)

    def process_image(self, image_path):
        image = Image.open(image_path)
        image = image.resize((800, 600))
        results = self.model.predict(source=image, conf=0.5)

        if len(results) == 0 or len(results[0].boxes.xywh) == 0:
            return image.resize((200, 150)), None

        first_object = results[0].boxes.xywh[0]
        x_center, y_center, width, height = first_object
        x1 = int((x_center - width / 2))
        y1 = int((y_center - height / 2))
        x2 = int((x_center + width / 2))
        y2 = int((y_center + height / 2))
        cropped_image = image.crop((x1, y1, x2, y2))
        cropped_image = cropped_image.resize((200, 150))
        return cropped_image, (x1, y1, x2, y2)

    def cancel_processing(self):
        self.cancel_flag = True

class FolderImageViewer(QMainWindow):
    def __init__(self, parent_folder):
        super().__init__()
        self.parent_folder = parent_folder
        self.child_folders = sorted(
            [
                os.path.join(parent_folder, d) 
                for d in os.listdir(parent_folder)
                if os.path.isdir(os.path.join(parent_folder, d)) and d.startswith("group")
            ],
            key=lambda x: int(re.search(r'group_(\d+)', x).group(1)) if re.search(r'group_(\d+)', x) else x
        )
        self.current_folder_index = 0
        self.image_files = []
        self.selected_images = set() 
        self.model = YOLO(MODEL_PATH)
        self.image_processor_thread = None
        self.image_bboxes = {}  # Store bounding boxes for each image

        self.init_ui()
        self.show_folder()

    def init_ui(self):
        self.setWindowTitle("Image Processing")
        self.setGeometry(100, 100, WINDOW_WIDTH, WINDOW_HEIGHT)

        # Main container
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)

        # Main vertical layout
        self.main_layout = QVBoxLayout()

        # Section 1: Vertically scrollable grid of images
        self.grid_scroll_area = QScrollArea()
        self.grid_widget = QWidget()
        self.grid_layout = QGridLayout()
        self.grid_widget.setLayout(self.grid_layout)
        self.grid_scroll_area.setWidget(self.grid_widget)
        self.grid_scroll_area.setWidgetResizable(True)
        self.grid_scroll_area.setFixedHeight(int(WINDOW_HEIGHT * .7))  # Adjust height as needed
        self.main_layout.addWidget(self.grid_scroll_area)

        # Section 2: Horizontally scrollable row of images
        self.row_scroll_area = QScrollArea()
        self.row_widget = QWidget()
        self.row_layout = QHBoxLayout()
        self.row_widget.setLayout(self.row_layout)
        self.row_scroll_area.setWidget(self.row_widget)
        self.row_scroll_area.setWidgetResizable(True)
        self.row_scroll_area.setFixedHeight(int(WINDOW_HEIGHT * .2))  # Adjust height as needed
        self.main_layout.addWidget(self.row_scroll_area)

        # Section 3: Controls (Buttons)
        rename_layout = QHBoxLayout()
        self.rename_input = QLineEdit(self)
        self.rename_input.setPlaceholderText("Enter new folder name...")
        rename_layout.addWidget(self.rename_input)

        self.rename_button = QPushButton("Apply Changes")
        self.rename_button.clicked.connect(self.rename_and_next)
        rename_layout.addWidget(self.rename_button)

        # Add the controls layout to the main layout
        self.main_layout.addLayout(rename_layout)

        # Set the central widget's layout
        self.central_widget.setLayout(self.main_layout)

    def show_folder(self):
        if not self.child_folders:
            print("No folders found.")
            return

        folder_path = self.child_folders[self.current_folder_index]
        self.image_files = [
            os.path.join(folder_path, f) for f in os.listdir(folder_path)
            if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))
        ]
        self.selected_images.clear()  # Clear selected images
        self.current_image_index = 0
        self.rename_input.setText("bib")
        self.image_bboxes.clear()  # Clear stored bounding boxes

        # Clear the grid layout
        while self.grid_layout.count():
            item = self.grid_layout.takeAt(0)
            widget = item.widget()
            if widget:
                widget.deleteLater()

        # Clear the grid layout
        while self.row_layout.count():
            item = self.row_layout.takeAt(0)
            widget = item.widget()
            if widget:
                widget.deleteLater()

        # Add images to the grid
        for i, image_file in enumerate(self.image_files):
            pixmap = QPixmap(image_file)
            IMAGE_SIZE = int((WINDOW_WIDTH-100) / COL_NUMS)
            pixmap = pixmap.scaled(IMAGE_SIZE, IMAGE_SIZE, Qt.KeepAspectRatio, Qt.SmoothTransformation)

            image_label = QLabel(self)
            image_label.setPixmap(pixmap)
            image_label.setAlignment(Qt.AlignCenter)
            image_label.setStyleSheet("border: 1px solid black;")
            image_label.image_path = image_file  # Store the image path
            image_label.mousePressEvent = lambda event, path=image_file, label=image_label: self.toggle_selection(event, path, label)

            row = i // COL_NUMS  # Assume 3 columns in the grid
            col = i % COL_NUMS
            self.grid_layout.addWidget(image_label, row, col)
        
        self.rename_input.setText("bib")
        self.crop_bibs_for_display()
    
    def toggle_selection(self, event, image_path, label):
        """Toggle selection of an image."""
        if image_path in self.selected_images:
            # Deselect: Remove from selected and reset styling
            self.selected_images.remove(image_path)
            label.setStyleSheet("border: 1px solid black;")
            
            # Remove the overlay if it exists
            if hasattr(label, 'overlay_label'):
                label.overlay_label.hide()
        else:
            # Select: Add to selected and apply styling
            self.selected_images.add(image_path)
            label.setStyleSheet("border: 2px solid red;")
            
            # Add an overlay with a red X
            if not hasattr(label, 'overlay_label'):
                overlay = QLabel(label)
                overlay.setGeometry(label.rect())
                overlay.setStyleSheet("background: rgba(255, 0, 0, 50%);")  # Semi-transparent red background
                overlay.setAlignment(Qt.AlignCenter)
                overlay.setText("X")
                overlay.setStyleSheet("color: red; font-size: 200px; font-weight: bold;")
                label.overlay_label = overlay
            label.overlay_label.show()

    def crop_bibs_for_display(self):
        if not self.image_files:
            return
        # Create and start the image processor thread
        self.image_processor_thread = ImageProcessorThread(self.image_files, self.model)  # Process all images
        self.image_processor_thread.image_processed.connect(self.add_image_to_grid)
        self.image_processor_thread.processing_canceled.connect(self.on_processing_canceled)
        self.image_processor_thread.start()
    
    def add_image_to_grid(self, idx, cropped_image, bbox):
        image_path = f"temp_image_{idx}.jpg"
        cropped_image.save(image_path)

        pixmap = QPixmap(image_path)
        image_label = QLabel(self)
        image_label.setPixmap(pixmap)
        image_label.setAlignment(Qt.AlignCenter)
        self.row_layout.addWidget(image_label)

        # Store the bounding box for the corresponding main image
        if idx < len(self.image_files):
            self.image_bboxes[self.image_files[idx]] = bbox
            # Add bounding box overlay to the main image
            self.add_bbox_overlay(self.image_files[idx], bbox)

    def add_bbox_overlay(self, image_path, bbox):
        if not bbox:
            return

        # Find the label widget for this image
        for i in range(self.grid_layout.count()):
            widget = self.grid_layout.itemAt(i).widget()
            if isinstance(widget, QLabel) and hasattr(widget, 'image_path') and widget.image_path == image_path:
                # Create a semi-transparent red rectangle overlay
                overlay = QLabel(widget)
                x1, y1, x2, y2 = bbox
                
                # Get the actual displayed image size and position
                pixmap = widget.pixmap()
                if not pixmap:
                    return
                    
                # Calculate scaling factors
                original_width = 800  # Original image width
                original_height = 600  # Original image height
                
                # Calculate the actual displayed image dimensions
                displayed_width = pixmap.width()
                displayed_height = pixmap.height()
                
                # Calculate the position of the image within the label
                label_width = widget.width()
                label_height = widget.height()
                x_offset = (label_width - displayed_width) // 2
                y_offset = (label_height - displayed_height) // 2
                
                # Scale the coordinates
                scale_x = displayed_width / original_width
                scale_y = displayed_height / original_height
                
                # Calculate the final coordinates
                final_x1 = int(x1 * scale_x) + x_offset
                final_y1 = int(y1 * scale_y) + y_offset
                final_width = int((x2 - x1) * scale_x)
                final_height = int((y2 - y1) * scale_y)
                
                overlay.setGeometry(final_x1, final_y1, final_width, final_height)
                overlay.setStyleSheet("background: rgba(255, 0, 0, 30%); border: 2px solid red;")
                overlay.show()
                break

    def on_processing_canceled(self):
        print("Image processing was canceled.")
        # self.clear_image_grid()
    
    def rename_and_next(self):
        new_name = self.rename_input.text().strip()
        if new_name:
            self.delete_selected_images()
            current_folder_path = self.child_folders[self.current_folder_index]
            parent_path = os.path.dirname(current_folder_path)
            new_folder_path = os.path.join(parent_path, new_name)

            if os.path.exists(new_folder_path):
                # Move all files from the current folder to the new folder
                for file_name in os.listdir(current_folder_path):
                    file_path = os.path.join(current_folder_path, file_name)
                    if os.path.isfile(file_path):  # Only move files, not subdirectories
                        os.rename(file_path, os.path.join(new_folder_path, file_name))
                os.rmdir(current_folder_path)  # Remove the now-empty folder
            else:
                os.rename(current_folder_path, new_folder_path)

            self.child_folders[self.current_folder_index] = new_folder_path
            print(f"Renamed folder to: {new_name}")
            self.show_next_folder()

        else:
            print("New folder name cannot be empty.")
    
    def show_next_folder(self):
        if self.child_folders:
            self.current_folder_index = (self.current_folder_index + 1)
            self.show_folder()
    
    def delete_selected_images(self):
        for image_path in self.selected_images:
            if os.path.exists(image_path):
                os.remove(image_path)
                print(f"Deleted: {image_path}")
        self.selected_images.clear()
    
    def keyPressEvent(self, event):
        if event.key() in {Qt.Key_Return, Qt.Key_Enter}:
            self.rename_and_next()
    
    def clearBothGrids(self):
        # Clear the grid layout
        while self.image_grid_layout.count():
            item = self.image_grid_layout.takeAt(0)
            widget = item.widget()
            if widget:
                widget.deleteLater()


if __name__ == "__main__":

    if SPLIT_INTO_GROUPS:
        imgSorter = SplitImagesByTimestamp(IMAGE_PATH, OUTPUT_PATH, TIME_DIFFERENCE)
        groups = imgSorter.group_photos_by_time()
        imgSorter.create_output_folders(groups)


        print(f"Photos have been grouped and saved in '{OUTPUT_PATH}'.")

    app = QApplication(sys.argv)

    if os.path.isdir(OUTPUT_PATH):
        viewer = FolderImageViewer(OUTPUT_PATH)
        viewer.show()
        sys.exit(app.exec_())
    else:
        print("Invalid folder path.")