BAT / main.py
aiyubali's picture
planogram updated
cb90fd0
# main.py
import json
import pandas as pd
from PIL import Image
import cv2
import numpy as np
from aiohttp import ClientSession
from io import BytesIO
import asyncio
import logging
from Data.config import *
from Utils.segment import CropModel
from Utils.sorting import SortModel
from Data.config import frameModel
from Data.config import blankModel
import matplotlib.pyplot as plt
# Set up logging
logging.basicConfig(level=logging.INFO)
class ImageFetcher:
@staticmethod
async def fetch_image(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:
logging.error(f"Failed to fetch image from {url}, status code: {response.status}")
return None
except Exception as e:
logging.error(f"Exception during image fetching from {url}: {e}")
return None
class ImageProcessor:
def __init__(self):
self.frameModel = frameModel
self.blankModel = blankModel
async def process_image(self, img_url):
async with ClientSession() as session:
image = await ImageFetcher.fetch_image(img_url, session)
if image is None:
return {"error": "Failed to fetch image"}
image = np.array(image) # Convert PIL image to NumPy array if needed
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (frame_shape[0], frame_shape[1])) # Resize to match the frame shape
try:
crop_model = CropModel(model=frameModel) # Initialize with your model
padded_image, other_class_name = crop_model.get_predicted_warped_image(image)
if other_class_name:
print(f"Other class name: {other_class_name}")
else:
print("No other class detected.")
# plt.imshow(padded_image)
# plt.axis('off') # Hide axes
# plt.show()
sort_model = SortModel(model=blankModel, classes_to_delete=classes_to_delete, conf_threshold=det_conf, expected_segments=expected_segments)
sequence_str = sort_model.process_image(padded_image, other_class_name)
return {"result": sequence_str}
except Exception as e:
logging.error(f"Exception during image processing: {e}")
return {"error": "Image processing failed."}