mirror of
https://github.com/mii443/usls.git
synced 2025-08-22 15:45:41 +00:00
Options -> ModelConfig
This commit is contained in:
@ -1,6 +1,6 @@
|
||||
use usls::{
|
||||
models::{TrOCR, TrOCRKind},
|
||||
DataLoader, Options, Scale,
|
||||
DataLoader, ModelConfig, Scale,
|
||||
};
|
||||
|
||||
#[derive(argh::FromArgs)]
|
||||
@ -38,52 +38,22 @@ fn main() -> anyhow::Result<()> {
|
||||
])?;
|
||||
|
||||
// build model
|
||||
let (options_encoder, options_decoder, options_decoder_merged) =
|
||||
match args.scale.as_str().try_into()? {
|
||||
Scale::S => match args.kind.as_str().try_into()? {
|
||||
TrOCRKind::Printed => (
|
||||
Options::trocr_encoder_small_printed(),
|
||||
Options::trocr_decoder_small_printed(),
|
||||
Options::trocr_decoder_merged_small_printed(),
|
||||
),
|
||||
TrOCRKind::HandWritten => (
|
||||
Options::trocr_encoder_small_handwritten(),
|
||||
Options::trocr_decoder_small_handwritten(),
|
||||
Options::trocr_decoder_merged_small_handwritten(),
|
||||
),
|
||||
},
|
||||
Scale::B => match args.kind.as_str().try_into()? {
|
||||
TrOCRKind::Printed => (
|
||||
Options::trocr_encoder_base_printed(),
|
||||
Options::trocr_decoder_base_printed(),
|
||||
Options::trocr_decoder_merged_base_printed(),
|
||||
),
|
||||
TrOCRKind::HandWritten => (
|
||||
Options::trocr_encoder_base_handwritten(),
|
||||
Options::trocr_decoder_base_handwritten(),
|
||||
Options::trocr_decoder_merged_base_handwritten(),
|
||||
),
|
||||
},
|
||||
x => anyhow::bail!("Unsupported TrOCR scale: {:?}", x),
|
||||
};
|
||||
let config = match args.scale.as_str().try_into()? {
|
||||
Scale::S => match args.kind.as_str().try_into()? {
|
||||
TrOCRKind::Printed => ModelConfig::trocr_small_printed(),
|
||||
TrOCRKind::HandWritten => ModelConfig::trocr_small_handwritten(),
|
||||
},
|
||||
Scale::B => match args.kind.as_str().try_into()? {
|
||||
TrOCRKind::Printed => ModelConfig::trocr_base_printed(),
|
||||
TrOCRKind::HandWritten => ModelConfig::trocr_base_handwritten(),
|
||||
},
|
||||
x => anyhow::bail!("Unsupported TrOCR scale: {:?}", x),
|
||||
}
|
||||
.with_device_all(args.device.as_str().try_into()?)
|
||||
.with_dtype_all(args.dtype.as_str().try_into()?)
|
||||
.commit()?;
|
||||
|
||||
let mut model = TrOCR::new(
|
||||
options_encoder
|
||||
.with_model_device(args.device.as_str().try_into()?)
|
||||
.with_model_dtype(args.dtype.as_str().try_into()?)
|
||||
.with_batch_size(xs.len())
|
||||
.commit()?,
|
||||
options_decoder
|
||||
.with_model_device(args.device.as_str().try_into()?)
|
||||
.with_model_dtype(args.dtype.as_str().try_into()?)
|
||||
.with_batch_size(xs.len())
|
||||
.commit()?,
|
||||
options_decoder_merged
|
||||
.with_model_device(args.device.as_str().try_into()?)
|
||||
.with_model_dtype(args.dtype.as_str().try_into()?)
|
||||
.with_batch_size(xs.len())
|
||||
.commit()?,
|
||||
)?;
|
||||
let mut model = TrOCR::new(config)?;
|
||||
|
||||
// inference
|
||||
let ys = model.forward(&xs)?;
|
||||
|
Reference in New Issue
Block a user