pineapple-neural 0.0.5

Simplified processing for image-based cell profiling with pineapple
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
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

// Copyright (c) 2025, Tom Ouellette
// Licensed under the BSD 3-Clause License

use candle_core::{DType, Device, Result};
use candle_nn::VarBuilder;

use pineapple_data::data::Weights;

use crate::models::DinoVisionTransformer;
use crate::models::{StandardVisionTransformer, StandardVisionTransformerConfig};

pub fn load_dinov2_vit_small(device: &Device, verbose: bool) -> Result<DinoVisionTransformer> {
    let weights = Weights::DinoVitSmall;
    let path = weights.path();
    weights.download(verbose);

    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[path], DType::F32, device)? };
    let model = DinoVisionTransformer::new(vb, 12, 384, 6, 14, 518).unwrap();

    Ok(model)
}

pub fn load_dinov2_vit_base(device: &Device, verbose: bool) -> Result<DinoVisionTransformer> {
    let weights = Weights::DinoVitBase;
    let path = weights.path();
    weights.download(verbose);

    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[path], DType::F32, device)? };
    let model = DinoVisionTransformer::new(vb, 12, 768, 12, 14, 518).unwrap();

    Ok(model)
}

pub fn load_dinobloom_vit_base(device: &Device, verbose: bool) -> Result<DinoVisionTransformer> {
    let weights = Weights::DinobloomVitBase;
    let path = weights.path();
    weights.download(verbose);

    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[path], DType::F32, device)? };
    let model = DinoVisionTransformer::new(vb, 12, 768, 12, 14, 224).unwrap();

    Ok(model)
}

pub fn load_subcell_vit_base(device: &Device, verbose: bool) -> Result<StandardVisionTransformer> {
    let weights = Weights::SubcellVitBase;
    let path = weights.path();
    weights.download(verbose);

    let config = StandardVisionTransformerConfig::vit_base_subcell();
    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[path], DType::F32, device)? };
    let model = StandardVisionTransformer::new(&config, vb).unwrap();

    Ok(model)
}

pub fn load_scdino_vit_small(device: &Device, verbose: bool) -> Result<StandardVisionTransformer> {
    let weights = Weights::ScdinoVitSmall;
    let path = weights.path();
    weights.download(verbose);

    let config = StandardVisionTransformerConfig::vit_base_scdino();
    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[path], DType::F32, device)? };
    let model = StandardVisionTransformer::new(&config, vb).unwrap();

    Ok(model)
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_load_dinov2_small() {
        let model = load_dinov2_vit_small(&Device::Cpu, true);
        assert!(model.is_ok());
    }

    #[test]
    fn test_load_dinov2_base_imagenet() {
        let model = load_dinov2_vit_base(&Device::Cpu, true);
        assert!(model.is_ok());
    }

    #[test]
    fn test_load_dinobloom_vit_base() {
        let model = load_dinobloom_vit_base(&Device::Cpu, true);
        assert!(model.is_ok());
    }

    #[test]
    fn test_load_subcell_vit_base() {
        let model = load_subcell_vit_base(&Device::Cpu, true);
        assert!(model.is_ok());
    }

    #[test]
    fn test_load_scdino_vit_small() {
        let model = load_scdino_vit_small(&Device::Cpu, true);
        assert!(model.is_ok());
    }
}