File size: 2,545 Bytes
f215795
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
import sys
from PIL import Image

# Append the path to the inference script's directory
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'src', 'slimface', 'inference')))
from end2end_inference import cinference_and_confirm

def run_inference(image, reference_dict_path, index_to_class_mapping_path, model_path, 

                 edgeface_model_name="edgeface_base", edgeface_model_dir="ckpts/idiap", 

                 algorithm="yolo", accelerator="auto", resolution=224, similarity_threshold=0.6):
    # Save uploaded image temporarily in apps/gradio_app/
    temp_image_path = os.path.join(os.path.dirname(__file__), "temp_image.jpg")
    image.save(temp_image_path)

    # Create args object to mimic command-line arguments
    class Args:
        def __init__(self):
            self.unknown_image_path = temp_image_path
            self.reference_dict_path = reference_dict_path.name if reference_dict_path else None
            self.index_to_class_mapping_path = index_to_class_mapping_path.name if index_to_class_mapping_path else None
            self.model_path = model_path.name if model_path else None
            self.edgeface_model_name = edgeface_model_name
            self.edgeface_model_dir = edgeface_model_dir
            self.algorithm = algorithm
            self.accelerator = accelerator
            self.resolution = resolution
            self.similarity_threshold = similarity_threshold

    args = Args()

    # Validate inputs
    if not all([args.reference_dict_path, args.index_to_class_mapping_path, args.model_path]):
        return "Error: Please provide all required files (reference dict, index-to-class mapping, and model)."

    try:
        # Call the inference function from end2end_inference.py
        results = cinference_and_confirm(args)
        
        # Format output
        output = ""
        for result in results:
            output += f"Image: {result['image_path']}\n"
            output += f"Predicted Class: {result['predicted_class']}\n"
            output += f"Confidence: {result['confidence']:.4f}\n"
            output += f"Similarity: {result.get('similarity', 'N/A'):.4f}\n"
            output += f"Confirmed: {result.get('confirmed', 'N/A')}\n\n"
        
        return output
    
    except Exception as e:
        return f"Error: {str(e)}"
    
    finally:
        # Clean up temporary image
        if os.path.exists(temp_image_path):
            os.remove(temp_image_path)