Spaces:
Runtime error
Runtime error
File size: 4,972 Bytes
22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 22d4865 63e95b2 | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | import json
import datetime
import httpx
import aioodbc
import numpy as np
from PIL import Image, ImageFile
from io import BytesIO
from fastapi import Request
from .BaseControllers import BaseControllers
from .FeatureExtractionControllers import FeatureExtractionControllers
from qdrant_client.http.models import PointStruct
class SQLDatabaseControllers(BaseControllers):
def __init__(self, request: Request, is_collections_recreated: bool):
super().__init__()
self.request = request
# Load last sync time
with open("last_sync_time.json", 'r') as f:
last_sync_data = json.load(f)
last_sync_str = last_sync_data.get("last_sync_time", datetime.datetime.min.isoformat())
self.last_sync_time = datetime.datetime.fromisoformat(last_sync_str)
if is_collections_recreated:
self.last_sync_time = datetime.datetime.min
self.feature_extractor = FeatureExtractionControllers(request=request)
ImageFile.LOAD_TRUNCATED_IMAGES = True
async def get_new_rows_query(self):
query = """
SELECT Id, PictureUrl, ColorId, Description
FROM Furniture
WHERE CreatedAt > ? OR UpdatedAt > ?
"""
async with self.request.app.SQLDatabasePool.acquire() as conn:
async with conn.cursor() as cursor:
await cursor.execute(query, (self.last_sync_time, self.last_sync_time))
rows = await cursor.fetchall()
columns = [column[0] for column in cursor.description]
return [dict(zip(columns, row)) for row in rows]
async def get_furniture_color(self, color_id: str):
query = """
SELECT Name
FROM Colors
WHERE Id = ?
"""
async with self.request.app.SQLDatabasePool.acquire() as conn:
async with conn.cursor() as cursor:
await cursor.execute(query, (color_id,))
row = await cursor.fetchone()
if row:
columns = [column[0] for column in cursor.description]
result = dict(zip(columns, row))
return result["Name"]
return "Unknown"
async def get_img_from_url(self, img_url: str):
if not img_url or img_url.strip() == "":
print("Empty or invalid image URL provided.")
return None
try:
async with httpx.AsyncClient() as client:
response = await client.get(img_url)
response.raise_for_status()
if "image" not in response.headers.get("Content-Type", ""):
print(f"Not an image: {img_url}")
return None
return Image.open(BytesIO(response.content)).convert("RGB")
except Exception as e:
print(f"Error fetching image: {e}")
return None
def extract_image_features(self, img: Image.Image):
return self.feature_extractor.extract_image_features(
image=img,
model=self.request.app.model,
processor=self.request.app.processor
)
def extract_text_features(self, description: str):
return self.feature_extractor.extract_text_features(
description=description,
model=self.request.app.model,
processor=self.request.app.processor
)
def batch_upsert(self, client, collection_name, points, batch_size=30):
for i in range(0, len(points), batch_size):
batch = points[i:i + batch_size]
client.upsert(collection_name=collection_name, points=batch)
async def sync_data(self):
rows = await self.get_new_rows_query()
if not rows:
return None
img_points = []
txt_points = []
for row in rows:
img = await self.get_img_from_url(row['PictureUrl'])
if not img:
continue
color = await self.get_furniture_color(row['ColorId'])
image_vector = self.extract_image_features(img)
image_vector = np.array(image_vector, dtype=np.float32).squeeze(axis=0)
img_points.append(PointStruct(
id=str(row['Id']),
vector=image_vector,
payload={"color": color}
))
text_vector = self.extract_text_features(row['Description'])
text_vector = np.array(text_vector, dtype=np.float32).squeeze(axis=0)
txt_points.append(PointStruct(
id=str(row['Id']),
vector=text_vector
))
self.batch_upsert(self.request.app.client, "text_features", txt_points)
self.batch_upsert(self.request.app.client, "image_features", img_points)
with open("last_sync_time.json", 'w') as f:
json.dump({
"last_sync_time": datetime.datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
}, f, indent=2)
|