Update handler.py

handler.py

@@ -12,104 +12,106 @@ from huggingface_hub import hf_hub_download
 from train_categories import PaddingImageProcessor
 
 def load_model_and_config(model_path):
-    """Carga el modelo entrenado y su configuración"""
-    hf_path = "vit_multiclass_model_best"
-    # Cargar información de las clases
-    class_info_path = os.path.join(hf_path, 'class_info.json')
-    with open(class_info_path, 'r') as f:
-        class_info = json.load(f)
-    
-    # Cargar configuración del procesador
-    processor_config_path = os.path.join(hf_path, 'processor_config.json')
-    with open(processor_config_path, 'r') as f:
-        processor_config = json.load(f)
-    
-    # Crear procesador de imágenes
-    image_processor = PaddingImageProcessor(
-        target_size=processor_config['target_size'],
-        padding_color=tuple(processor_config['padding_color'])
-    )
-    
-    # Cargar modelo
-    model = ViTForImageClassification.from_pretrained(model_path)
-    model.eval()
-    
-    # Usar GPU si está disponible
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = model.to(device)
-    
-    return model, image_processor, class_info, device
+	"""Carga el modelo entrenado y su configuración"""
+	hf_path = "vit_multiclass_model_best"
+	# Cargar información de las clases
+	class_info_path = os.path.join(model_path, 'class_info.json')
+	with open(class_info_path, 'r') as f:
+		class_info = json.load(f)
+	
+	# Cargar configuración del procesador
+	processor_config_path = os.path.join(model_path, 'processor_config.json')
+	with open(processor_config_path, 'r') as f:
+		processor_config = json.load(f)
+	
+	# Crear procesador de imágenes
+	image_processor = PaddingImageProcessor(
+		target_size=processor_config['target_size'],
+		padding_color=tuple(processor_config['padding_color'])
+	)
+	
+	# Cargar modelo
+	model = ViTForImageClassification.from_pretrained(model_path)
+	model.eval()
+	
+	# Usar GPU si está disponible
+	device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+	model = model.to(device)
+	
+	return model, image_processor, class_info, device
 
 def download_image(url: str) -> Image.Image:
-    """Descarga una imagen desde una URL"""
-    response = requests.get(url, timeout=10)
-    response.raise_for_status()
-    image = Image.open(BytesIO(response.content)).convert('RGB')
-    return image
+	"""Descarga una imagen desde una URL"""
+	response = requests.get(url, timeout=10)
+	response.raise_for_status()
+	image = Image.open(BytesIO(response.content)).convert('RGB')
+	return image
 
 def classify_image(model, image_processor, class_info, device, accuracy):
 
-    # Descargar y procesar imagen
-    image = download_image(image_url)
-    processed_image = image_processor(image).unsqueeze(0).to(device)
-    
-    # Realizar predicción
-    with torch.no_grad():
-        outputs = model(pixel_values=processed_image).logits
-        probabilities = torch.sigmoid(outputs).cpu().numpy()[0]
-    
-    # Obtener clases predichas (umbral 0.5)
-    predicted_classes = []
-    for i, prob in enumerate(probabilities):
-        if prob > accuracy:
-            class_name = class_info['class_columns'][i]
-            predicted_classes.append(f"{class_name}: {prob:.3f}")
-    
-    # Mostrar resultado
-    if predicted_classes:
-        for prediction in predicted_classes:
-            print(prediction)
-        return predicted_classes    
-    else:
-        # Si ninguna clase supera 0.5, mostrar la más probable
-        max_idx = probabilities.argmax()
-        max_prob = probabilities[max_idx]
-        class_name = class_info['class_columns'][max_idx]
-        print(f"{class_name}: {max_prob:.3f}")
-        return [class_name, max_prob]
+	# Descargar y procesar imagen
+	image = download_image(image_url)
+	processed_image = image_processor(image).unsqueeze(0).to(device)
+	
+	# Realizar predicción
+	with torch.no_grad():
+		outputs = model(pixel_values=processed_image).logits
+		probabilities = torch.sigmoid(outputs).cpu().numpy()[0]
+	
+	# Obtener clases predichas (umbral 0.5)
+	predicted_classes = []
+	for i, prob in enumerate(probabilities):
+		if prob > accuracy:
+			class_name = class_info['class_columns'][i]
+			predicted_classes.append(f"{class_name}: {prob:.3f}")
+	
+	# Mostrar resultado
+	if predicted_classes:
+		for prediction in predicted_classes:
+			print(prediction)
+		return predicted_classes    
+	else:
+		# Si ninguna clase supera 0.5, mostrar la más probable
+		max_idx = probabilities.argmax()
+		max_prob = probabilities[max_idx]
+		class_name = class_info['class_columns'][max_idx]
+		print(f"{class_name}: {max_prob:.3f}")
+		return [class_name, max_prob]
 
 class EndpointHandler():
-    def __init__(self, path=""):
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        model_filename = "vit_multiclass_model_best/model.safetensors"
-        local_path = hf_hub_download(repo_id="Drazcat-AI/categories_peru", filename=model_filename)
-        print('MODEL PATH:',local_path)
-        self.model, self.image_processor, self.class_info, self.device = load_model_and_config(local_path)
+	def __init__(self, path=""):
+		device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+		model_filename = "vit_multiclass_model_best/model.safetensors"
+		local_path = hf_hub_download(repo_id="Drazcat-AI/categories_peru", filename=model_filename)
+		hf_hub_download(repo_id="Drazcat-AI/categories_peru", filename='class_info.json')
+		hf_hub_download(repo_id="Drazcat-AI/categories_peru", filename='processor_config.json')
+		print('MODEL PATH:',local_path)
+		self.model, self.image_processor, self.class_info, self.device = load_model_and_config(local_path)
 
-    def predict_objects(self, image_url, accuracy):
-        
-        result_df =  classify_image(image_url, accuracy)   
-        return result_df
+	def predict_objects(self, image_url, accuracy):
+		
+		result_df =  classify_image(image_url, accuracy)   
+		return result_df
 
-    def __call__(self, event):
-        if "inputs" not in event:
-            return {
-                "statusCode": 400,
-                "body": json.dumps("Error: Please provide an 'inputs' parameter."),
-            }
-        event = event["inputs"]
-        image_url = event["image_url"]
-        accuracy = event["accuracy"]
-        try:
-            predictions = self.predict_objects(image_url, accuracy)
-            predictions_json = predictions.to_json(orient='records')
-            
-            return {
-                "statusCode": 200,
-                "body": json.dumps(predictions_json),
-            }
-        except Exception as e:
-            return {
-                "statusCode": 500,
-                "body": json.dumps(f"Error: {str(e)}"),
-            }      
+	def __call__(self, event):
+		if "inputs" not in event:
+			return {
+				"statusCode": 400,
+				"body": json.dumps("Error: Please provide an 'inputs' parameter."),
+			}
+		event = event["inputs"]
+		image_url = event["image_url"]
+		accuracy = event["accuracy"]
+		try:
+			predictions = self.predict_objects(image_url, accuracy)
+			predictions_json = predictions.to_json(orient='records')
+			
+			return {
+				"statusCode": 200,
+				"body": json.dumps(predictions_json),
+			}
+		except Exception as e:
+			return {
+				"statusCode": 500,
+				"body": json.dumps(f"Error: {str(e)}"),
+			}