ImageTools.py

from rembg import remove
from PIL import Image, ImageTk
import os
import torch
from ultralytics import YOLO
import cv2 as cv
import numpy as np
import tkinter as tk
import matplotlib.pyplot as plt

class Scratch:
    def __init__(self, image1_path, image2_path = 'aligned_img.jpg'):
        self.gd_img = cv.imread(image1_path)
        self.align_img = cv.imread(image2_path)
        self.r = self.gd_img.shape[0]//2 # for floodfill
    def crop_inner_circle(self, img, radius_reduction_factor=0.85):
        gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        blurred = cv.GaussianBlur(gray, (9, 9), 0)
        _, binary = cv.threshold(blurred, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)
        contours, _ = cv.findContours(binary, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
        # do the trick
        max_contour = max(contours, key=cv.contourArea)
        (x, y), radius = cv.minEnclosingCircle(max_contour)
        # Reduce the radius to 95 perc
        inner_radius = int(radius * radius_reduction_factor)
        center = (int(x), int(y))
    
        mask = np.zeros_like(gray)
        cv.circle(mask, center, inner_radius, 255, -1)
        result = cv.bitwise_and(img, img, mask=mask)
        # Crop the bounding rectangle of the inner circle
        x1, y1 = max(0, center[0] - inner_radius), max(0, center[1] - inner_radius)
        x2, y2 = min(img.shape[1], center[0] + inner_radius), min(img.shape[0], center[1] + inner_radius)
        #cropped = result[y1:y2, x1:x2]
        #cv.imwrite(output_path, cropped)
        return x1, y1, x2, y2
    def circle_crop(self, image_path):
        image = cv.imread(image_path)
        if image is None:
            raise ValueError("Image not found or path is invalid.")
        height, width = image.shape[:2]
        radius = min(width, height) // 2
        center_x, center_y = width // 2, height // 2
        center = (center_x, center_y)
        mask = np.zeros((height, width), dtype=np.uint8)
        cv.circle(mask, center, radius, 255, -1)
        result = cv.bitwise_and(image, image, mask=mask)
        x1, y1 = center_x - radius, center_y - radius
        x2, y2 = center_x + radius, center_y + radius
        cropped_result = result[y1:y2, x1:x2]
        return cropped_result



''' 少用的2個功能 '''
class rgb_yolo:
    def rbg(self, input_dir, output_dir): # 花很久, 至少30秒
        os.makedirs(output_dir, exist_ok=True)
        for file_name in os.listdir(input_dir):
            input_path = os.path.join(input_dir, file_name)
            output_path = os.path.join(output_dir, file_name)
            if file_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                try:
                    with open(input_path, 'rb') as input_file:
                        img_data = input_file.read()
                        result = remove(img_data)  
                    with open(output_path, 'wb') as output_file:
                        output_file.write(result)  
                except Exception as e:
                    print(f"Error processing {file_name}: {e}")
    def run_yolo(self, input_dir, output_dir, confidence=0.25): # 蠻快的
        model = YOLO('yolov8l.pt')
        model.conf = confidence
        os.makedirs(output_dir, exist_ok=True)
        for file_name in os.listdir(input_dir):
            input_path = os.path.join(input_dir, file_name)
            output_path = os.path.join(output_dir, file_name)
            img = cv.imread(input_path)
            if file_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                try:
                    results = model(input_path)
                    high_conf_results = [r for r in results[0].boxes if r.conf.item() >= 0.1 and r.cls.item()==74.0]
                    if high_conf_results:
                        for i, r in enumerate(high_conf_results):
                            x1, y1, x2, y2 = map(int, r.xyxy[0].tolist())
                            high_conf_img = img[y1:y2, x1:x2]
                            cv.imwrite(output_path, high_conf_img)
                    else:
                        print(f"No objects detected in {file_name}. Skipping.")
                except Exception as e:
                    print(f"Error processing {file_name}: {e}")