Skip to content

Commit

Permalink
Customizable naming for ImageExtractor
Browse files Browse the repository at this point in the history 
Summary: This diff isolates the naming of extracted files in a customizable class, so it's easy to reuse ImageExtractor with a different file naming scheme.

Reviewed By: davidhahnfb, kiminoue7

Differential Revision: D51444583

fbshipit-source-id: 8c133aa037f63d54c42763e3b166aef636f279a7
  • Loading branch information
@facebook-github-bot
Georges Berenger authored and facebook-github-bot committed Nov 27, 2023
1 parent e5ba39b commit fdf4399
Show file tree
Hide file tree
Showing 2 changed files with 109 additions and 54 deletions.
111 changes: 60 additions & 51 deletions vrs/utils/ImageExtractor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,17 +29,27 @@ using namespace std;
using namespace vrs;
using namespace utils;

namespace {
namespace vrs {
namespace utils {

const bool kSupportGrey16Export = true;
string ImageNamer::namePngImage(StreamId id, uint32_t imageCounter, double timestamp) {
return fmt::format("{}-{:05}-{:.3f}.png", id.getNumericName(), imageCounter, timestamp);
}

bool writeRawImage(
string ImageNamer::nameRawImage(
const ImageContentBlockSpec& imageSpec,
const vector<uint8_t>& imageData,
const string& folderPath,
StreamId id,
uint32_t imageCounter,
double timestamp) {
return fmt::format(
"{}-{:05}-{:.3f}{}",
id.getNumericName(),
imageCounter,
timestamp,
getRawImageFormatAsString(imageSpec));
}

string ImageNamer::getRawImageFormatAsString(const ImageContentBlockSpec& imageSpec) {
const auto& imageFormat = imageSpec.getImageFormat();
string filenamePostfix;
string extension;
Expand Down Expand Up @@ -68,16 +78,17 @@ bool writeRawImage(
extension = toString(imageFormat);
break;
}
return fmt::format("{}.{}", filenamePostfix, extension);
}

string path = fmt::format(
"{}/{}-{:05}-{:.3f}{}.{}",
folderPath,
id.getNumericName(),
imageCounter,
timestamp,
filenamePostfix,
extension);
} // namespace utils
} // namespace vrs

namespace {

const bool kSupportGrey16Export = true;

bool writeRawImage(const string& path, const vector<uint8_t>& imageData) {
ofstream file(path, ios::binary);
if (!file.is_open()) {
XR_LOGE("Cannot open file {} for writing", path);
Expand All @@ -95,30 +106,20 @@ bool writeRawImage(
enum JobType { SaveToPng, EndQueue };

struct ImageJob {
ImageJob(const string& folderPath, StreamId id, double timestamp, uint32_t imageCounter)
: jobType{JobType::SaveToPng},
folderPath{folderPath},
id{id},
timestamp{timestamp},
imageCounter{imageCounter},
frame{make_shared<PixelFrame>()} {}
ImageJob(JobType jobType, const string& folderPath) : jobType{jobType}, folderPath{folderPath} {}
ImageJob(const string&& path)
: jobType{JobType::SaveToPng}, path{path}, frame{make_shared<PixelFrame>()} {}
ImageJob() : jobType{EndQueue} {}

void saveAsPng() {
shared_ptr<PixelFrame> normalFrame;
PixelFrame::normalizeFrame(frame, normalFrame, kSupportGrey16Export);
string path = fmt::format(
"{}/{}-{:05}-{:.3f}.png", folderPath, id.getNumericName(), imageCounter, timestamp);
fmt::print("Writing {}\n", path);
normalFrame->writeAsPng(path);
}

JobType jobType;

const string& folderPath;
StreamId id;
double timestamp;
uint32_t imageCounter;
const string path;
shared_ptr<PixelFrame> frame;
};

Expand All @@ -142,10 +143,8 @@ class ImageProcessor {

void endThreadPool() {
unique_lock<mutex> locker(mutex_);
if (threadPool_.size() > 0) {
JobQueue<unique_ptr<ImageJob>>& imageQueue = getImageQueue();
string fakePath;
imageQueue.sendJob(make_unique<ImageJob>(JobType::EndQueue, fakePath));
if (!threadPool_.empty()) {
imageQueue_.sendJob(make_unique<ImageJob>());
for (auto& thread : threadPool_) {
if (thread.joinable()) {
thread.join();
Expand Down Expand Up @@ -176,24 +175,40 @@ class ImageProcessor {
JobQueue<unique_ptr<ImageJob>> imageQueue_;
};

ImageNamer& getDefaultImageNamer() {
static ImageNamer sDefaultImageNamer;
return sDefaultImageNamer;
}

} // namespace

namespace vrs {
namespace utils {

ImageExtractor::ImageExtractor(const string& folderPath, uint32_t& counter, bool extractImagesRaw)
: ImageExtractor(getDefaultImageNamer(), folderPath, counter, extractImagesRaw) {}

ImageExtractor::ImageExtractor(
ImageNamer& imageNamer,
const string& folderPath,
uint32_t& counter,
const bool extractImagesRaw)
: folderPath_{folderPath}, imageFileCounter_{counter}, extractImagesRaw_(extractImagesRaw) {
bool extractImagesRaw)
: imageNamer_{imageNamer},
folderPath_{folderPath},
imageFileCounter_{counter},
extractImagesRaw_(extractImagesRaw) {
ImageProcessor::get().startThreadPool();
}

ImageExtractor::~ImageExtractor() {
ImageProcessor::get().endThreadPool();
}

bool ImageExtractor::onImageRead(const CurrentRecord& record, size_t, const ContentBlock& ib) {
bool ImageExtractor::onDataLayoutRead(const CurrentRecord& r, size_t idx, DataLayout& dl) {
return imageNamer_.onDataLayoutRead(r, idx, dl, *this);
}

bool ImageExtractor::onImageRead(const CurrentRecord& record, size_t idx, const ContentBlock& ib) {
JobQueue<unique_ptr<ImageJob>>& imageQueue = ImageProcessor::get().getImageQueue();
while (imageQueue.getQueueSize() > 2 * thread::hardware_concurrency()) {
this_thread::sleep_for(chrono::milliseconds(50));
Expand All @@ -204,17 +219,12 @@ bool ImageExtractor::onImageRead(const CurrentRecord& record, size_t, const Cont
const StreamId id = record.streamId;
auto format = ib.image().getImageFormat();

if (!extractImagesRaw_ && format == ImageFormat::RAW) {
unique_ptr<ImageJob> job =
make_unique<ImageJob>(folderPath_, id, record.timestamp, imageCounter_);
if (PixelFrame::readRawFrame(job->frame, record.reader, ib.image())) {
imageQueue.sendJob(std::move(job));
return true;
}
} else if (!extractImagesRaw_ && format == ImageFormat::VIDEO) {
unique_ptr<ImageJob> job =
make_unique<ImageJob>(folderPath_, id, record.timestamp, imageCounter_);
if (tryToDecodeFrame(*job->frame, record, ib) == 0) {
if (!extractImagesRaw_ && (format == ImageFormat::RAW || format == ImageFormat::VIDEO)) {
string filename = imageNamer_.namePngImage(id, imageCounter_, record.timestamp);
unique_ptr<ImageJob> job = make_unique<ImageJob>(fmt::format("{}/{}", folderPath_, filename));
if ((format == ImageFormat::RAW &&
PixelFrame::readRawFrame(job->frame, record.reader, ib.image())) ||
(format == ImageFormat::VIDEO && tryToDecodeFrame(*job->frame, record, ib) == 0)) {
imageQueue.sendJob(std::move(job));
return true;
}
Expand All @@ -231,21 +241,20 @@ bool ImageExtractor::onImageRead(const CurrentRecord& record, size_t, const Cont
errorCodeToMessage(readStatus));
return false;
}
writeRawImage(ib.image(), imageData, folderPath_, id, imageCounter_, record.timestamp);
string filename = imageNamer_.nameRawImage(ib.image(), id, imageCounter_, record.timestamp);
string filepath = fmt::format("{}/{}", folderPath_, filename);
writeRawImage(filepath, imageData);
return true;
}
XR_LOGW("Could not convert image for {}, format: {}", id.getName(), ib.asString());
return false;
}

bool ImageExtractor::onUnsupportedBlock(
const CurrentRecord& record,
size_t,
const ContentBlock& cb) {
bool ImageExtractor::onUnsupportedBlock(const CurrentRecord& r, size_t i, const ContentBlock& cb) {
// the image was not decoded, probably because the image spec are incomplete
if (cb.getContentType() == ContentType::IMAGE) {
imageCounter_++;
XR_LOGW("Image skipped for {}, content: {}", record.streamId.getName(), cb.asString());
XR_LOGW("Image skipped for {}, content: {}", r.streamId.getName(), cb.asString());
}
return false;
}
Expand Down
52 changes: 49 additions & 3 deletions vrs/utils/ImageExtractor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,17 +26,63 @@
namespace vrs {
namespace utils {

class ImageExtractor;

/// Optional helper class so ImageExtractor's image naming can be customized
class ImageNamer {
public:
virtual ~ImageNamer() = default;

/// Before reading any record, after the file is open, use this callback to know what is read.
virtual void init(RecordFileReader& reader) {}

/// For each record in image streams, get their datalayouts.
/// Use extractor.getExpectedLayout<T>() as needed.
virtual bool onDataLayoutRead(
const CurrentRecord& record,
size_t blockIndex,
DataLayout& datalayout,
ImageExtractor& extractor) {
return true;
}

/// Name image files saved in png format (default)
virtual string namePngImage(StreamId id, uint32_t imageCounter, double timestamp);

/// Name image files saved in raw format
virtual string nameRawImage(
const ImageContentBlockSpec& imageSpec,
StreamId id,
uint32_t imageCounter,
double timestamp);

protected:
string getRawImageFormatAsString(const ImageContentBlockSpec& imageSpec);
};

class ImageExtractor : public utils::VideoRecordFormatStreamPlayer {
public:
ImageExtractor(const string& folderPath, uint32_t& counter, const bool extractImagesRaw);
ImageExtractor(const string& folderPath, uint32_t& counter, bool extractImagesRaw);
ImageExtractor(
ImageNamer& imageNamer,
const string& folderPath,
uint32_t& counter,
bool extractImagesRaw);
~ImageExtractor() override;

bool onImageRead(const CurrentRecord& record, size_t, const ContentBlock& ib) override;
bool onUnsupportedBlock(const CurrentRecord& record, size_t, const ContentBlock& cb) override;
bool onDataLayoutRead(const CurrentRecord& r, size_t idx, DataLayout& dl) override;
bool onImageRead(const CurrentRecord& record, size_t idx, const ContentBlock& ib) override;
bool onUnsupportedBlock(const CurrentRecord& record, size_t idx, const ContentBlock& cb) override;

void saveImagesThreadActivity();

template <class T>
inline T& getExpectedLayout(DataLayout& layout, size_t blockIndex) {
return RecordFormatStreamPlayer::getExpectedLayout<T>(layout, blockIndex);
}

protected:
ImageNamer& imageNamer_;
const string& folderPath_;
uint32_t& imageFileCounter_;
uint32_t imageCounter_ = 0;
Expand Down

0 comments on commit fdf4399

Please sign in to comment.