torsh-vision 0.1.2

Computer vision utilities for ToRSh deep learning framework
Documentation 
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

use crate::{ModelConfig, Result, VisionError, VisionModel};
use std::collections::HashMap;

/// Model registry for vision models
pub struct ModelRegistry {
    models: HashMap<String, String>,
}

impl ModelRegistry {
    /// Create a new model registry
    pub fn new() -> Self {
        let mut registry = Self {
            models: HashMap::new(),
        };
        registry.register_default_models();
        registry
    }

    /// Register default models
    fn register_default_models(&mut self) {
        self.models
            .insert("resnet18".to_string(), "ResNet-18".to_string());
        self.models
            .insert("resnet34".to_string(), "ResNet-34".to_string());
        self.models
            .insert("resnet50".to_string(), "ResNet-50".to_string());
        self.models
            .insert("alexnet".to_string(), "AlexNet".to_string());
        self.models
            .insert("vgg16".to_string(), "VGG-16".to_string());
    }

    /// List available models
    pub fn list_models(&self) -> Vec<String> {
        self.models.keys().cloned().collect()
    }

    /// Check if model exists
    pub fn has_model(&self, name: &str) -> bool {
        self.models.contains_key(name)
    }
}

impl Default for ModelRegistry {
    fn default() -> Self {
        Self::new()
    }
}