mirror of
https://github.com/mii443/usls.git
synced 2025-08-22 15:45:41 +00:00
Add florence2 model
* Add florence2-base model for all tasks * Update annotator.rs
This commit is contained in:
252
examples/florence2/main.rs
Normal file
252
examples/florence2/main.rs
Normal file
@ -0,0 +1,252 @@
|
||||
use usls::{models::Florence2, Annotator, DataLoader, Options, Task};
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// vision encoder
|
||||
let options_vision_encoder = Options::default()
|
||||
.with_model("florence2/base-vision-encoder.onnx")?
|
||||
.with_i00((1, 2, 4).into())
|
||||
.with_i02((512, 768, 800).into())
|
||||
.with_i03((512, 768, 800).into())
|
||||
.with_profile(false)
|
||||
.with_cuda(0);
|
||||
|
||||
// text embed
|
||||
let options_text_embed = Options::default()
|
||||
.with_model("florence2/base-embed-tokens.onnx")?
|
||||
.with_i00((1, 2, 4).into())
|
||||
.with_i01((1, 2, 20).into()) // seq_length
|
||||
.with_tokenizer("florence2/tokenizer.json")?
|
||||
.with_profile(false);
|
||||
|
||||
// transformer encoder
|
||||
let options_encoder = Options::default()
|
||||
.with_model("florence2/base-encoder.onnx")?
|
||||
.with_i00((1, 2, 4).into())
|
||||
.with_i01((1, 2, 300).into()) // encoder_sequence_length
|
||||
.with_i10((1, 2, 4).into())
|
||||
.with_i11((1, 2, 300).into()) // encoder_sequence_length
|
||||
.with_profile(false);
|
||||
|
||||
// transformer decoder
|
||||
let options_decoder = Options::default()
|
||||
.with_model("florence2/base-decoder.onnx")?
|
||||
.with_i00((1, 2, 4).into())
|
||||
.with_i01((1, 2, 300).into()) // encoder_sequence_length
|
||||
.with_i10((1, 2, 4).into())
|
||||
.with_i11((1, 2, 300).into()) // encoder_sequence_length
|
||||
.with_i20((1, 2, 4).into())
|
||||
.with_i21((1, 2, 300).into()) // encoder_sequence_length
|
||||
.with_profile(false);
|
||||
|
||||
// transformer decoder merged
|
||||
let options_decoder_merged = Options::default()
|
||||
.with_model("florence2/base-decoder-merged.onnx")?
|
||||
// encoder_attention_mask
|
||||
.with_i00((1, 2, 4).into())
|
||||
.with_i01((1, 2, 300).into()) // encoder_sequence_length
|
||||
// encoder_hidden_states
|
||||
.with_i10((1, 2, 4).into())
|
||||
.with_i11((1, 2, 300).into()) // encoder_sequence_length
|
||||
// inputs_embeds
|
||||
.with_i20((1, 2, 4).into())
|
||||
.with_i21((1, 2, 300).into()) // encoder_sequence_length
|
||||
// past_key_values.0.decoder.key
|
||||
.with_i30((1, 2, 4).into())
|
||||
.with_i32_((1, 2, 1).into())
|
||||
// past_key_values.0.decoder.value
|
||||
.with_i40((1, 2, 4).into())
|
||||
.with_i42((1, 2, 1).into())
|
||||
// past_key_values.0.encoder.key
|
||||
.with_i50((1, 2, 4).into())
|
||||
.with_i52((1, 2, 1).into())
|
||||
// past_key_values.0.decoder.value
|
||||
.with_i60((1, 2, 4).into())
|
||||
.with_i62((1, 2, 1).into())
|
||||
// past_key_values.1.decoder.key
|
||||
.with_i70((1, 2, 4).into())
|
||||
.with_i72((1, 2, 1).into())
|
||||
// past_key_values.1.decoder.value
|
||||
.with_i80((1, 2, 4).into())
|
||||
.with_i82((1, 2, 1).into())
|
||||
// past_key_values.1.encoder.key
|
||||
.with_i90((1, 2, 4).into())
|
||||
.with_i92((1, 2, 1).into())
|
||||
// past_key_values.1.decoder.value
|
||||
.with_i100((1, 2, 4).into())
|
||||
.with_i102((1, 2, 1).into())
|
||||
// past_key_values.2.decoder.key
|
||||
.with_i110((1, 2, 4).into())
|
||||
.with_i112((1, 2, 1).into())
|
||||
// past_key_values.2.decoder.value
|
||||
.with_i120((1, 2, 4).into())
|
||||
.with_i122((1, 2, 1).into())
|
||||
// past_key_values.2.encoder.key
|
||||
.with_i130((1, 2, 4).into())
|
||||
.with_i132((1, 2, 1).into())
|
||||
// past_key_values.2.decoder.value
|
||||
.with_i140((1, 2, 4).into())
|
||||
.with_i142((1, 2, 1).into())
|
||||
// past_key_values.3.decoder.key
|
||||
.with_i150((1, 2, 4).into())
|
||||
.with_i152((1, 2, 1).into())
|
||||
// past_key_values.3.decoder.value
|
||||
.with_i160((1, 2, 4).into())
|
||||
.with_i162((1, 2, 1).into())
|
||||
// past_key_values.3.encoder.key
|
||||
.with_i170((1, 2, 4).into())
|
||||
.with_i172((1, 2, 1).into())
|
||||
// past_key_values.3.decoder.value
|
||||
.with_i180((1, 2, 4).into())
|
||||
.with_i182((1, 2, 1).into())
|
||||
// past_key_values.4.decoder.key
|
||||
.with_i190((1, 2, 4).into())
|
||||
.with_i192((1, 2, 1).into())
|
||||
// past_key_values.4.decoder.value
|
||||
.with_i200((1, 2, 4).into())
|
||||
.with_i202((1, 2, 1).into())
|
||||
// past_key_values.4.encoder.key
|
||||
.with_i210((1, 2, 4).into())
|
||||
.with_i212((1, 2, 1).into())
|
||||
// past_key_values.4.decoder.value
|
||||
.with_i220((1, 2, 4).into())
|
||||
.with_i222((1, 2, 1).into())
|
||||
// past_key_values.5.decoder.key
|
||||
.with_i230((1, 2, 4).into())
|
||||
.with_i232((1, 2, 1).into())
|
||||
// past_key_values.5.decoder.value
|
||||
.with_i240((1, 2, 4).into())
|
||||
.with_i242((1, 2, 1).into())
|
||||
// past_key_values.5.encoder.key
|
||||
.with_i250((1, 2, 4).into())
|
||||
.with_i252((1, 2, 1).into())
|
||||
// past_key_values.5.decoder.value
|
||||
.with_i260((1, 2, 4).into())
|
||||
.with_i262((1, 2, 1).into())
|
||||
//use_cache_branch
|
||||
.with_i270((1, 2, 1).into())
|
||||
.with_profile(false);
|
||||
|
||||
// build model
|
||||
let mut model = Florence2::new(
|
||||
options_vision_encoder,
|
||||
options_text_embed,
|
||||
options_encoder,
|
||||
options_decoder,
|
||||
options_decoder_merged,
|
||||
)?;
|
||||
|
||||
// load images
|
||||
let xs = [
|
||||
// DataLoader::try_read("florence2/car.jpg")?, // for testing region-related tasks
|
||||
DataLoader::try_read("florence2/car.jpg")?,
|
||||
// DataLoader::try_read("images/db.png")?,
|
||||
DataLoader::try_read("assets/bus.jpg")?,
|
||||
];
|
||||
|
||||
// region-related tasks
|
||||
let quantizer = usls::Quantizer::default();
|
||||
// let coords = [449., 270., 556., 372.]; // wheel
|
||||
let coords = [31., 156., 581., 373.]; // car
|
||||
let (width_car, height_car) = (xs[0].width(), xs[0].height());
|
||||
let quantized_coords = quantizer.quantize(&coords, (width_car as _, height_car as _));
|
||||
|
||||
// run with tasks
|
||||
let ys = model.run_with_tasks(
|
||||
&xs,
|
||||
&[
|
||||
// w/ inputs
|
||||
Task::Caption(0),
|
||||
Task::Caption(1),
|
||||
Task::Caption(2),
|
||||
Task::Ocr,
|
||||
Task::OcrWithRegion,
|
||||
Task::RegionProposal,
|
||||
Task::ObjectDetection,
|
||||
Task::DenseRegionCaption,
|
||||
// w/o inputs
|
||||
Task::OpenSetDetection("a vehicle".into()),
|
||||
Task::CaptionToPhraseGrounding(
|
||||
"A vehicle with two wheels parked in front of a building.".into(),
|
||||
),
|
||||
Task::ReferringExpressionSegmentation("a vehicle".into()),
|
||||
Task::RegionToSegmentation(
|
||||
quantized_coords[0],
|
||||
quantized_coords[1],
|
||||
quantized_coords[2],
|
||||
quantized_coords[3],
|
||||
),
|
||||
Task::RegionToCategory(
|
||||
quantized_coords[0],
|
||||
quantized_coords[1],
|
||||
quantized_coords[2],
|
||||
quantized_coords[3],
|
||||
),
|
||||
Task::RegionToDescription(
|
||||
quantized_coords[0],
|
||||
quantized_coords[1],
|
||||
quantized_coords[2],
|
||||
quantized_coords[3],
|
||||
),
|
||||
],
|
||||
)?;
|
||||
|
||||
// annotator
|
||||
let annotator = Annotator::new()
|
||||
.without_bboxes_conf(true)
|
||||
.with_bboxes_thickness(3)
|
||||
.with_saveout_subs(&["Florence2"]);
|
||||
for (task, ys_) in ys.iter() {
|
||||
match task {
|
||||
Task::Caption(_)
|
||||
| Task::Ocr
|
||||
| Task::RegionToCategory(..)
|
||||
| Task::RegionToDescription(..) => {
|
||||
println!("Task: {:?}\n{:?}\n", task, ys_)
|
||||
}
|
||||
Task::DenseRegionCaption => {
|
||||
let annotator = annotator.clone().with_saveout("Dense-Region-Caption");
|
||||
annotator.annotate(&xs, ys_);
|
||||
}
|
||||
Task::RegionProposal => {
|
||||
let annotator = annotator
|
||||
.clone()
|
||||
.without_bboxes_name(false)
|
||||
.with_saveout("Region-Proposal");
|
||||
|
||||
annotator.annotate(&xs, ys_);
|
||||
}
|
||||
Task::ObjectDetection => {
|
||||
let annotator = annotator.clone().with_saveout("Object-Detection");
|
||||
annotator.annotate(&xs, ys_);
|
||||
}
|
||||
Task::OpenSetDetection(_) => {
|
||||
let annotator = annotator.clone().with_saveout("Open-Set-Detection");
|
||||
annotator.annotate(&xs, ys_);
|
||||
}
|
||||
Task::CaptionToPhraseGrounding(_) => {
|
||||
let annotator = annotator
|
||||
.clone()
|
||||
.with_saveout("Caption-To-Phrase-Grounding");
|
||||
annotator.annotate(&xs, ys_);
|
||||
}
|
||||
Task::ReferringExpressionSegmentation(_) => {
|
||||
let annotator = annotator
|
||||
.clone()
|
||||
.with_saveout("Referring-Expression-Segmentation");
|
||||
annotator.annotate(&xs, ys_);
|
||||
}
|
||||
Task::RegionToSegmentation(..) => {
|
||||
let annotator = annotator.clone().with_saveout("Region-To-Segmentation");
|
||||
annotator.annotate(&xs, ys_);
|
||||
}
|
||||
Task::OcrWithRegion => {
|
||||
let annotator = annotator.clone().with_saveout("Ocr-With-Region");
|
||||
annotator.annotate(&xs, ys_);
|
||||
}
|
||||
|
||||
_ => (),
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
Reference in New Issue
Block a user