main.py

import json
import pandas as pd
from PIL import Image
from aiohttp import ClientSession
from io import BytesIO
import asyncio
from Data.model import fridgeModel, drinksModel
from Data.data import pepsi_items, competitor_items, water_items  

class ImageFetcher:
    async def fetch_image(self, url, session):
        try:
            async with session.get(url) as response:
                if response.status == 200:
                    img_data = await response.read()
                    return Image.open(BytesIO(img_data))
                else:
                    print(f"Failed to fetch image from {url}, status code: {response.status}")
                    return None
        except Exception as e:
            print(f"Exception during image fetching from {url}: {e}")
            return None

class DetectionFilter:
    @staticmethod
    def filter_detection(detection_dict, category_list):
        filtered = {}
        for name, count in detection_dict.items():
            if name in category_list:
                filtered[name] = count
        return filtered

class ImageDetector:
    def __init__(self, model, thresh):
        self.model = model
        self.thresh = thresh

    async def detect_items(self, urls, session):
        detection = {}
        fetcher = ImageFetcher()
        try:
            for url in urls:
                image = await fetcher.fetch_image(url, session)
                if image:
                    results = self.model(image, conf=self.thresh)
                    if len(results) > 0:
                        data = json.loads(results[0].tojson())
                        df = pd.DataFrame(data)
                        #print("Dataframe:", df)
                        if 'name' in df.columns:
                            name_counts = df['name'].value_counts().sort_index()
                            for name, count in name_counts.items():
                                if name in detection:
                                    detection[name] += count
                                else:
                                    detection[name] = count
                        else:
                            print(f"No 'name' column found in the DataFrame for URL: {url}")
                    else:
                        print(f"No results found for image from URL: {url}")
                else:
                    print(f"No image fetched for URL: {url}")
        except Exception as e:
            print(f"Error during detection: {e}")
        return detection

class ImageProcessor:
    def __init__(self):
        # Initialize models (Category lists are now imported directly)
        self.fridge_model = fridgeModel
        self.drinks_model = drinksModel

    async def process_images(self, fdz_urls, citem_urls):
        async with ClientSession() as session:
            # Run detection tasks concurrently for both models
            fridge_detector = ImageDetector(self.fridge_model, thresh=0.8)
            drinks_detector = ImageDetector(self.drinks_model, thresh=0.6)

            fdz_detection = await fridge_detector.detect_items(fdz_urls, session)
            citem_detection = await drinks_detector.detect_items(citem_urls, session)

            # Filter citem_detection into categories
            filter_tool = DetectionFilter()
            pepsi = filter_tool.filter_detection(citem_detection, pepsi_items)
            competitor = filter_tool.filter_detection(citem_detection, competitor_items)
            water = filter_tool.filter_detection(citem_detection, water_items)

            # Construct skuDetection dictionary only if it has items
            sku_detection = {}
            if pepsi:
                sku_detection["pepsico"] = pepsi
            if competitor:
                sku_detection["competitor"] = competitor
            if water:
                sku_detection["water"] = water

            # Prepare response
            response = {
                "fdzDetection": fdz_detection,
                "skuDetection": sku_detection
            }

            return response