File size: 4,073 Bytes
7f4d7c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
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
|